#include <cstddef>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <numeric>
#include <vector>
#include<fstream>
#include<iostream>
#include<tuple>

#include "common/tensorrt.hpp"
#include "common/check.hpp"
#include "cuda_util.h"
#include "backbone_neck.hpp"


namespace dsgn{
    class BackboneImpl : public Backbone{
        public:
            virtual ~BackboneImpl() {
                if (left_stereo_feat_) checkRuntime(cudaFree(left_stereo_feat_));
                if (right_stereo_feat_) checkRuntime(cudaFree(right_stereo_feat_));
                if (left_sem_feat_) checkRuntime(cudaFree(left_sem_feat_));
                // if (right_sem_feat_) checkRuntime(cudaFree(right_sem_feat_));
                if (left_sem_feat_for_lsv_) checkRuntime(cudaFree(left_sem_feat_for_lsv_));
            }
            
            bool init(const std::string& model, const int& width, const int& height) {
                engine_ = TensorRT::load(model);
                if (engine_ == nullptr) return false;

                stereo_feat_dims_ = engine_->static_dims(2);
                sem_feat_dims_ = engine_->static_dims(3);
                sem_feat_for_lsv_shape_ = engine_->static_dims(5);
                size_t bytes_image_ = width * height * 3 * sizeof(nvtype::half);
                int32_t volumn = std::accumulate(stereo_feat_dims_.begin(), stereo_feat_dims_.end(), 1, std::multiplies<int32_t>());
                checkRuntime(cudaMalloc(&left_stereo_feat_, volumn * sizeof(nvtype::half)));
                checkRuntime(cudaMalloc(&right_stereo_feat_, volumn * sizeof(nvtype::half)));
                volumn = std::accumulate(sem_feat_dims_.begin(), sem_feat_dims_.end(), 1, std::multiplies<int32_t>());
                checkRuntime(cudaMalloc(&left_sem_feat_, volumn * sizeof(nvtype::half)));
                volumn = std::accumulate(sem_feat_for_lsv_shape_.begin(), sem_feat_for_lsv_shape_.end(), 1, std::multiplies<int32_t>());
                checkRuntime(cudaMalloc(&left_sem_feat_for_lsv_, volumn * sizeof(nvtype::half)));
                    return true;
                }

            virtual void print() override { engine_->print("Feature Extractor Backbone"); }

            virtual void forward(const nvtype::half* images, void* stream=nullptr) override {
                const nvtype::half* image_r = images+bytes_image_;
                engine_->forward({images, image_r,  
                                 left_stereo_feat_, right_stereo_feat_, 
                                 left_sem_feat_, left_sem_feat_for_lsv_},
                                 static_cast<cudaStream_t>(stream));
            }

            virtual nvtype::half* left_stereo_feat() override { return left_stereo_feat_; }
            virtual nvtype::half* right_stereo_feat() override { return right_stereo_feat_; }
            virtual nvtype::half* left_sem_feat() override { return left_sem_feat_; }
            virtual nvtype::half* left_sem_feat_for_lsv() override { return left_sem_feat_for_lsv_; }
            virtual std::vector<int> stereo_feat_shape() override { return stereo_feat_dims_; }
            virtual std::vector<int> sem_feat_shape() override { return sem_feat_dims_; }
            virtual std::vector<int> sem_feat_for_lsv_shape() override { return sem_feat_for_lsv_shape_; }

        private:
            std::shared_ptr<TensorRT::Engine> engine_;
            size_t bytes_image_;
            nvtype::half* left_stereo_feat_ = nullptr;
            nvtype::half* right_stereo_feat_ = nullptr;
            nvtype::half* left_sem_feat_ = nullptr;
            // nvtype::half* right_sem_feat_ = nullptr;
            nvtype::half* left_sem_feat_for_lsv_ = nullptr;
            std::vector<int> stereo_feat_dims_, sem_feat_dims_, sem_feat_for_lsv_shape_;
    };

    std::shared_ptr<Backbone> create_backbone(const std::string& model, const int& width, const int& height) {
        std::shared_ptr<BackboneImpl> instance(new BackboneImpl());
        if (!instance->init(model, width, height)){
            instance.reset();
        }
        return instance;
    };
} // namespace dsgn
